from sklearn.svm import SVC  # Import the Support Vector Classifier from scikit-learn
from sklearn.metrics import classification_report  # Import the classification report for evaluating the classifier
from tfidf import TFIDF  # Import the TFIDF class for text feature extraction

class DocumentClassifier:
    def __init__(self):
        # Initialize the TFIDF processor
        self.tfidf = TFIDF()
        # Initialize the Support Vector Classifier with a linear kernel
        self.classifier = SVC(kernel='linear')
        # Flag to check if the classifier has been trained
        self.is_trained = False
        
    def add_training_documents(self, documents, labels):
        """添加训练文档和对应的标签"""
        # Iterate over each document to add it to the TFIDF processor
        for doc in documents:
            self.tfidf.add_document(doc)
            
        # Generate feature vectors for each document
        X = [self.tfidf.get_feature_vector(doc) for doc in documents]
        # Assign labels to the documents
        y = labels
        
        # Train the classifier with the feature vectors and labels
        self.classifier.fit(X, y)
        # Set the trained flag to True
        self.is_trained = True
        
    def predict(self, document):
        """预测文档的类别"""
        # Check if the classifier is trained
        if not self.is_trained:
            raise ValueError("分类器尚未训练")
            
        # Generate the feature vector for the input document
        feature_vector = self.tfidf.get_feature_vector(document)
        
        # Predict the class of the document and return it
        return self.classifier.predict([feature_vector])[0]
    
    def evaluate(self, test_documents, test_labels):
        """评估分类器性能"""
        # Check if the classifier is trained
        if not self.is_trained:
            raise ValueError("分类器尚未训练")
            
        # Generate feature vectors for the test documents
        X_test = [self.tfidf.get_feature_vector(doc) for doc in test_documents]
        
        # Predict the classes for the test documents
        predictions = self.classifier.predict(X_test)
        
        # Generate and return the classification report
        return classification_report(test_labels, predictions)